import matplotlib.pyplot as plt
from pyscf import gto, scf

import sys
sys.path.append("../Hartree-Fock/")
from hartreefock import *

def benchmark_H2():
    r = np.linspace(1, 2)
    Emy = []
    Epyscf = []
    for ri in r:
        Emy.append(HartreeFock(H2(ri)))
        h2 = gto.M(atom=f"H 0 0 0; H 0 0 {ri*a0}", basis="sto-3g")
        Epyscf.append(scf.RHF(h2).kernel())
    
    plt.figure()
    plt.plot(r, Emy, '+', label="MY CODE")
    plt.plot(r, Epyscf, label="PySCF")
    plt.xlabel(r'$r/\mathrm{a.u.}$')
    plt.ylabel(r'$E/\mathrm{a.u.}$')
    plt.legend()
    plt.savefig("benchmark-H2.png")

def benchmark_HeH():
    r = np.linspace(1, 5)
    Emy = []
    Epyscf = []
    for ri in r:
        Emy.append(HartreeFock(HeH(ri)))
        heh = gto.M(atom=f"He 0 0 0; H 0 0 {ri*a0}", basis="sto-3g", charge=1)
        Epyscf.append(scf.RHF(heh).kernel())
    
    plt.figure()
    plt.plot(r, Emy, '+', label="MY CODE")
    plt.plot(r, Epyscf, label="PySCF")
    plt.xlabel(r'$r/\mathrm{a.u.}$')
    plt.ylabel(r'$E/\mathrm{a.u.}$')
    plt.legend()
    plt.savefig("benchmark-HeH.png")

if __name__ == '__main__':
    benchmark_H2()
    benchmark_HeH()